import torch
import numpy as np
import open3d as o3d
            
import matplotlib.pyplot as plt
from pytorch3d.structures import Meshes
from pytorch3d.renderer import FoVPerspectiveCameras
from pytorch3d.renderer import RasterizationSettings
from pytorch3d.renderer import MeshRenderer
from pytorch3d.renderer import MeshRasterizer
from pytorch3d.renderer import SoftPhongShader
from pytorch3d.renderer import PointLights
from pytorch3d.renderer import TexturesVertex

meshed_file = "/data/workspace/3D/room_ws/dense_291_individual/meshed-poisson.ply"
# 加载 ply 网格文件
mesh_o3d = o3d.io.read_triangle_mesh(meshed_file)
# 计算顶点的法线(normal)
mesh_o3d.compute_vertex_normals()
'''
获取网格化后的点（顶点）的张量
在pytorch中，无论是标量、向量还是多维数组，统称为张量（tensor）

unsqueeze（解压，squeeze是挤压之意），实现了增加了顶级维度（第0个维度），比如对于原始维度[N,3]，经过了unsqueeze之后，
就是变成了[1,N,3]，增加了一个顶级维度；之所以这样的操作，是因为基于张量tensor的很多操作都是会基于batch（批量）
这样顶级维度来做的，所以如果原始数据中没有batch顶级维度，就默认创建一个统一顶级维度编号为1的的维度，
用于匹配计算需要的维度
'''
verts = torch.tensor(np.asarray(mesh_o3d.vertices), dtype=torch.float32).unsqueeze(0)
# 获取网格化后的面（三角形）的张量
faces = torch.tensor(np.asarray(mesh_o3d.triangles), dtype=torch.int64).unsqueeze(0)

# 如果网格有颜色信息，则沿用网格颜色
if mesh_o3d.has_vertex_colors():
    colors = np.asarray(mesh_o3d.vertex_colors)
# 如果没有颜色信息，给默认白色（ones_like代表构建全1数组，在RGB里面全0对应黑色，全1代表白色）
else:
    colors = np.ones_like(np.asarray(mesh_o3d.vertices)) * 0.8

# 为颜色添加一个batch维度
textures = torch.tensor(colors, dtype=torch.float32).unsqueeze(0)

# 构建 Meshes 对象（pytorch3d的对象），参数分别对应open3d获取的顶点信息、三角面信息以及颜色信息
# 不过都是基于原始数据添加了一个batch维度（tensor计算维度要求的）
mesh = Meshes(verts=verts, faces=faces, textures=TexturesVertex(textures))

# 摄像机参数
device = "cuda:0" if torch.cuda.is_available() else "cpu"
print("device is: " + device)

# 设置计算设备类型，GPU/ CPU
device = torch.device(device)
mesh = mesh.to(device)

def render_image(camera_position, look_at, image_name="render.png"):
    # 存储一组参数，该组参数可以用于生成图像
    '''
    FoVPerspectiveCameras，模拟了一个真实世界中，带视场角的相机，用于透视效果的3D渲染图象
    重要参数说明：
    R：相机的旋转矩阵，控制相机的朝向，如果R是单位矩阵（对角线是1，其他位置是0），则说明没有旋转
    T：相机的位置（平移向量），控制相机在那里
    znear/ zfar：近远裁剪面，默认一般不用动
    fov：视野角度（默认60°），fov是Field of View（视场角）的缩写，对于宽FoV，虽然拍的多，但是会有畸变，而且Fov越大，畸变越大
    aspect_ratio：宽高比，默认是1.0（正方向）
    
    类似的还有OrthographicCameras，两者区别如下：
        =============================================
        类别	                 模拟
        =============================================
        FoVPerspectiveCameras	真实世界相机（近大远小）
        OrthographicCameras	    平行投影（没有透视）
        =============================================
    '''
    cameras = FoVPerspectiveCameras(device=device, R=look_at, T=camera_position)
    '''
    PointLights，点光源，用于模拟现实中的点光源，像小灯泡一样向四周发光
    构造函数中设置了一组默认的光源颜色（有默认颜色），可以保持默认光源颜色，但是，需要指定光源的位置，
    光源的位置一般设置为摄像机的位置，或者比摄像机略高：
    1. ambient_color：环境光
    2. diffuse_color：漫反射
    3. specular_color：（镜面）反射光
    '''
    lights = PointLights(device=device, location=[[0.0, 0.0, -3.0]])
    
    # raster：光栅，raster_settings，光栅设置
    '''
    RasterizationSettings定义了如何将网格（三角形）投影到2D图像的规则；
    它通常在创建MeshRasterizer的时候，用于设置渲染尺寸、面片采样方式、抗锯齿精度等。
    ==================================================================================================
    参数名	                类型	含义
    ==================================================================================================
    image_size	            int     或 tuple(int, int)	输出图像的尺寸（宽高）。可以是整数或 (H, W) 元组
    blur_radius	            float	模糊半径。用于抗锯齿或者软渲染，设置为 0 表示无模糊（硬边）
    faces_per_pixel         int	    每个像素最多保留多少个面片的信息（默认1），用于软渲染、阴影等效果，如果
                                    大于1，代表开启软光栅化，能实现半透明、边缘平滑以及阴影模拟等效果
    bin_size	            int	    光栅化加速的分块大小（可选，一般不需设置）
    max_faces_per_bin	    int	    每个块最多支持的面片数量，超过则报错
    perspective_correct	    bool	是否进行透视校正插值（默认 True）
    clip_barycentric_coords	bool	是否裁剪重心坐标，避免出界（默认 True）
    ==================================================================================================
    '''
    raster_settings = RasterizationSettings(
        image_size=512, # wide & height的像素数
        blur_radius=0.0,
        faces_per_pixel=1,
    )
    '''
    ===========================================================================
    Shader	                特点
    ===========================================================================
    SoftPhongShader	        模拟光照（最常用），有立体感，图像尽量接近“现实感”
    SoftSilhouetteShader	只画轮廓（用于 mask 任务）
    HardPhongShader	        更硬朗的光照，不模糊边界
    SoftGouraudShader	    光照在顶点计算，再插值到像素（比 Phong 快但不如细腻）
    ===========================================================================
    '''
    renderer = MeshRenderer(
        rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
        shader=SoftPhongShader(device=device, cameras=cameras, lights=lights),
    )

    images = renderer(mesh)
    image = images[0, ..., :3].cpu().numpy()

    plt.imsave(image_name, image)
    print(f"Saved: {image_name}")

'''
设置摄像机的位置和朝向, camera_pos设定Z轴坐标为-3.0，说明当前点在Z轴的负方向，然后at变量是旋转，而通过
torch.eys获得的矩阵是单位矩阵，单位矩阵代表相对于原点没有旋转，即look at指向的就是正前方，这样

    camera_pos+loop at

的组合，代表面朝向原点，即Z+方向，对于[[0.0, 0.0, -3.0]，Z轴坐标-3，如果换成-100，则相当于拉远摄像头，类似的
如何换成-1，相当于拉近摄像机：

          ↑ Y+
          |
          | 
          |_____→ X+
         /
        /
      Z+

类似，如果通过下面的get_rotation_matrix函数指定了旋转矩阵，则代表面向Z+方向，旋转指定角度

'''
camera_pos = torch.tensor([[0.0, 0.0, -3.0]], device=device)  # 摄像机在原点前方
# 朝向固定为 Z+
# 只要是张量，一般都是需要在需要将正常数据添加一个维度0的batch维度
at = torch.eye(3, device=device).unsqueeze(0)  

# 渲染图像
render_image(camera_pos, at, image_name="camera_forward.png")

def get_rotation_matrix(yaw_deg):
    yaw = np.radians(yaw_deg)
    R = np.array([
        [np.cos(yaw), 0, np.sin(yaw)],
        [0, 1, 0],
        [-np.sin(yaw), 0, np.cos(yaw)]
    ])
    return torch.tensor(R, dtype=torch.float32)

# 模拟右转 30 度
right_rot = get_rotation_matrix(-30).unsqueeze(0).to(device)
render_image(camera_pos, right_rot, image_name="camera_right30.png")

# 模拟左转 30 度
left_rot = get_rotation_matrix(30).unsqueeze(0).to(device)
render_image(camera_pos, left_rot, image_name="camera_left30.png")
